
import torch
from torch import nn
from torchvision import transforms, models, datasets
from PIL import Image
from torchvision import transforms, datasets

def set_parameter_requires_grad(model, feature_extracting):
    if feature_extracting:
        for param in model.parameters():
            param.requires_grad = False

def initialize_model(model_name, num_classes, feature_extract, use_pretrained=True):
    # 选择合适的模型，不同模型的初始化方法稍微有点区别
    model_ft = None
    input_size = 0

    if model_name == "resnet":
        """ Resnet152
        """
        model_ft = models.resnet152(pretrained=use_pretrained)
        set_parameter_requires_grad(model_ft, feature_extract)
        num_ftrs = model_ft.fc.in_features
        model_ft.fc = nn.Sequential(nn.Linear(num_ftrs, 102),
                                   nn.LogSoftmax(dim=1))
        input_size = 224


    else:
        print("Invalid model name, exiting...")
        exit()

    return model_ft, input_size

def get_model():
    model_name = 'resnet'
    feature_extract = True
    model_ft, input_size = initialize_model(model_name, 2, feature_extract, use_pretrained=True)
    return model_ft

def predict_image(model,image_path):

    filename = './weights/bound_check.pt'
    checkpoint = torch.load(filename)
    best_acc = checkpoint['best_acc']
    class_names = ["0", "1"]
    model.load_state_dict(checkpoint['state_dict'])
    #print(checkpoint["state_dict"])
    model.eval()
    # （1）此处为使用PIL进行测试的代码
    transform_valid = transforms.Compose([
        transforms.Resize((56, 56), interpolation=2),
        transforms.ToTensor()
    ]
    )
    img = Image.open(image_path)
    img_ = transform_valid(img).unsqueeze(0)  # 拓展维度
    outputs = model(img_)
    _, indices = torch.max(outputs, 1)
    result = class_names[indices]
    #print('predicted:', result)
    return result
if __name__ == '__main__':
    model=get_model()
    predict_image(model,"1.jpg")



"""
import torch
from torch import nn
from PIL import Image
from torchvision import transforms,datasets 
import cv2

class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()       

        self.conv1 = nn.Conv2d(in_channels = 3, out_channels = 32, padding = 1,kernel_size = 3)
        self.pool1 = nn.MaxPool2d(kernel_size = 2,stride = 2)
        self.conv2 = nn.Conv2d(in_channels = 32, out_channels = 64, padding = 1,kernel_size = 3)

        self.conv3 = nn.Conv2d(in_channels = 64, out_channels = 128, padding = 1,kernel_size = 3)
        self.conv4 = nn.Conv2d(in_channels = 128, out_channels = 128, padding = 1, kernel_size = 3)

        self.relu = nn.ReLU()
        self.flatten = nn.Flatten()
        self.linear1 = nn.Linear(128 * 7 * 7, 512)
        self.linear2 = nn.Linear(512, 6)
        self.softmax = nn.Softmax()

    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.pool1(x)
        x = self.relu(self.conv2(x))
        x = self.pool1(x)

        x = self.relu(self.conv3(x))
        x = self.pool1(x)
        x = self.relu(self.conv4(x))

        x = self.flatten(x)
        x = self.relu(self.linear1(x))
        y = self.linear2(x)
        return y  

class_names = ['GC', 'GL', 'NL', 'RC', 'RL', 'UK'] #这个顺序很重要，要和训练时候的类名顺序一致

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

##载入模型并读取权重
model = Net()
model.load_state_dict(torch.load("./data/detect_light.pt"))
model.to(device)
model.eval()

img_path = '/home/jwd/dataset/roi455.jpg'

#（1）此处为使用PIL进行测试的代码
transform_valid = transforms.Compose([
    transforms.Resize((56, 56), interpolation=2),
    transforms.ToTensor()
  ]
)
img = Image.open(img_path)
img_ = transform_valid(img).unsqueeze(0) #拓展维度

##（2）此处为使用opencv读取图像的测试代码，若使用opencv进行读取，将上面（1）注释掉即可。
# img = cv2.imread(img_path)
# img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# img = cv2.resize(img, (56, 56))
# img_ = torch.from_numpy(img).float().permute(2, 0, 1).unsqueeze(0)/255

img_ = img_.to(device)
outputs = model(img_)

#输出概率最大的类别
_, indices = torch.max(outputs,1)
percentage = torch.nn.functional.softmax(outputs, dim=1)[0] * 100
perc = percentage[int(indices)].item()
result = class_names[indices]
print('predicted:', result)

# 得到预测结果，并且从大到小排序
# _, indices = torch.sort(outputs, descending=True)
# 返回每个预测值的百分数
# percentage = torch.nn.functional.softmax(outputs, dim=1)[0] * 100
# print([(class_names[idx], percentage[idx].item()) for idx in indices[0][:5]])
"""